import networkx as nx
import numpy as np
import glob
import os, os.path
import math
import pickle
import scipy.sparse as sp
import random

# pathhack = os.path.dirname(os.path.realpath(__file__))

# network = nx.Graph()

#处理网格lattice数据
import grid

FILE_NAME = 'grid_2d'
RANDOM_SEED = None
def load_edges():
    global network
    edge_file = open("./original/{}.txt".format(FILE_NAME),"r")
    for line in edge_file:
        # nodefrom nodeto
        split = [x for x in line.split("	")]
        node_from = int(split[0])
        node_to = int(split[1])
        network.add_edge(node_from, node_to)


if __name__ == '__main__':
    # print "Running tests."
    # print("Loading network...")
    # load_edges()
    # print("done.")
    #
    # failures = 0
    #
    # def test(actual, expected, test_name):
    #     global failures  # lol python scope
    #     try:
    #         print("testing %s..." % (test_name,))
    #         assert actual == expected, "%s failed (%s != %s)!" % (test_name, actual, expected)
    #         print("%s passed (%s == %s)." % (test_name, actual, expected))
    #     except AssertionError as e:
    #         print(e)
    #         failures += 1
    #
    length_row = 20
    length_col = 20
    data = grid.grid_2d(length_row,length_col)
    # G = nx.gnm_random_graph(2500,50000)
    # G = nx.gnp_random_graph(2500,0.5)
    # p_new_connection = 0.5
    # new_edges = []
    # for i in G.nodes():
    #     # find the other nodes this one is connected to
    #     connected = [to for (fr, to) in G.edges(node)]
    #     # and find the remainder of nodes, which are candidates for new edges
    #     unconnected = [n for n in G.nodes() if not n in connected]
    #
    #     # probabilistically add a random edge
    #     if len(unconnected):  # only try if new edge is possible
    #         if random.random() < p_new_connection:
    #             new = random.choice(unconnected)
    #             dis = abs(new[0] - node[0]) + abs(new[1] - node[1])
    #             if (dis>=length/3):
    #                 G.add_edge(node, new)
    #             print
    #             "\tnew edge:\t {} -- {}".format(node, new)
    #             new_edges.append((node, new))
    #             # book-keeping, in case both add and remove done in same cycle
    #             unconnected.remove(new)
    #             connected.append(new)
    #             np.random.seed(RANDOM_SEED)
    G,news_edges = data
    print('network size', G.size())
    # print("%d tests failed." % (failures,))
    print('')
    print("Saving network...")
    adj = nx.adjacency_matrix(G)
    res = adj,news_edges
    # adj_ = adj.toarray()
    print(adj.shape[0])

    file = 'D:/data-processed/{}-{}-{}-adj.pkl'.format(FILE_NAME,length_row,length_col)
    # os.makedir(file)
    with open(file, "wb") as f:
        pickle.dump(res, f)

    print("Network saved!")

